package com.fe.class08;

import java.util.HashMap;

/**
 * @Description 前缀树
 * 1）单个字符串中，字符从前到后的加到一棵多叉树上
 * 2）字符放在路上，节点上有专属的数据项（常见的是pass和end值）
 * 3）所有样本都这样添加，如果没有路就新建，如有路就复用
 * 4）沿途节点的pass值增加1，每个字符串结束时来到的节点end值增加1
 *
 * 功能：1）添加某个字符串，可以重复添加，每次算1个
 * 2）查询某个字符串在结构中还有几个
 * 3）查询有多少个字符串，是以str做前缀的
 * 4）删掉某个字符串，可以重复删除，每次算1个
 * @Author sosuke :-)
 * @Date 2022/1/10 22:45
 */
public class Code02_TrieTree {
    // 表示前缀树的节点
    public static class Node {
        private int pass;
        private int end;
        // key表示分支（字符对应的ASCII码）
        private HashMap<Integer, Node> nexts;

        public Node() {
            this.pass = 0;
            this.end = 0;
            this.nexts = new HashMap<>();
        }
    }

    // 表示前缀树结构
    public static class Trie {
        private Node root; // 根节点

        public Trie() {
            this.root = new Node();
        }

        /**
         * 插入字符串到前缀树中
         * 注意：字符放在分支路上
         *
         * @param word
         */
        public void insert(String word) {
            if (null == word || word.trim().length() == 0) return;
            char[] chars = word.toCharArray();
            Node node = root; // 从根节点开始
            node.pass++;
            // 遍历字符串的每个字符
            for (int i = 0; i < chars.length; i++) {
                int path = chars[i];
                if (!node.nexts.containsKey(path)) {
                    node.nexts.put(path, new Node());
                }
                node = node.nexts.get(path);
                node.pass++; // 沿途节点的pass值增加1
            }
            node.end++; // 每个字符串结束时来到的节点end值增加1
        }

        /**
         * 在前缀树中查询某字符串加入过几次
         *
         * @param word
         * @return
         */
        public int search(String word) {
            if (null == word || word.trim().length() == 0) return 0;
            char[] chars = word.toCharArray();
            Node node = root;
            for (int i = 0; i < chars.length; i++) {
                int path = chars[i];
                if (!node.nexts.containsKey(path)) {
                    return 0;
                }
                node = node.nexts.get(path);
            }
            return node.end; // 每个字符串结束时来到的节点end值增加1，所以这里返回node.end
        }

        /**
         * 在前缀树中查询有几个是以该字符串为前缀的
         *
         * @param word
         * @return
         */
        public int prefixNum(String word) {
            if (null == word || word.trim().length() == 0) return 0;
            char[] chars = word.toCharArray();
            Node node = root;
            for (int i = 0; i < chars.length; i++) {
                int path = chars[i];
                if (!node.nexts.containsKey(path)) {
                    return 0;
                }
                node = node.nexts.get(path);
            }
            return node.pass; // 沿途节点的pass值增加1，所以这里返回node.pass
        }

        /**
         * 在前缀树中删除某个字符串
         *
         * @param word
         */
        public void delete(String word) {
            if (search(word) != 0) { // 存在
                char[] chars = word.toCharArray();
                Node node = root;
                root.pass--;
                for (int i = 0; i < chars.length; i++) {
                    int path = chars[i];
                    if (--node.nexts.get(path).pass == 0) {
                        node.nexts.remove(path); // 释放节点，防止内存泄漏
                        return;
                    }
                    node = node.nexts.get(path);
                }
                node.end--;
            }
        }
    }

    // 使用hash表来实现前缀树的功能
    public static class TrieTest {
        // 使用hash表作为词频统计表
        private HashMap<String, Integer> box;

        public TrieTest() {
            box = new HashMap<>();
        }

        public void insert(String word) {
            if (!box.containsKey(word)) {
                box.put(word, 1);
            } else {
                box.put(word, box.get(word) + 1);
            }
        }

        public int search(String word) {
            return box.getOrDefault(word, 0);
        }

        public int prefixNumber(String pre) {
            int count = 0;
            for (String cur : box.keySet()) {
                if (cur.startsWith(pre)) {
                    count += box.get(cur);
                }
            }
            return count;
        }

        public void delete(String word) {
            if (box.containsKey(word)) {
                if (box.get(word) == 1) {
                    box.remove(word);
                } else {
                    box.put(word, box.get(word) - 1);
                }
            }
        }
    }

    public static void main(String[] args) {
        int testTimes = 100000;
        int arrLen = 100;
        int strLen = 10;
        for (int i = 0; i < testTimes; i++) {
            String[] arr = generateRandomStringArray(arrLen, strLen);
            Trie trie = new Trie();
            TrieTest trieTest = new TrieTest();
            for (int j = 0; j < arr.length; j++) {
                double decide = Math.random();
                if (decide < 0.25) {
                    trie.insert(arr[j]);
                    trieTest.insert(arr[j]);
                } else if (decide < 0.5) {
                    trie.delete(arr[j]);
                    trieTest.delete(arr[j]);
                } else if (decide < 0.75) {
                    int ans1 = trie.search(arr[j]);
                    int ans2 = trieTest.search(arr[j]);
                    if (ans1 != ans2) {
                        System.out.println("Oops!");
                    }
                } else {
                    int ans1 = trie.prefixNum(arr[j]);
                    int ans2 = trieTest.prefixNumber(arr[j]);
                    if (ans1 != ans2) {
                        System.out.println("Oops!");
                    }
                }
            }
        }
        System.out.println("finish!");
    }

    private static String[] generateRandomStringArray(int maxSize, int strLen) {
        int size = (int) (Math.random() * maxSize) + 1;
        String[] arr = new String[size];
        for (int i = 0; i < size; i++) {
            arr[i] = generateRandomString(strLen);
        }
        return arr;
    }

    private static String generateRandomString(int strLen) {
        int size = (int) (Math.random() * strLen) + 1;
        char[] chars = new char[size];
        for (int i = 0; i < chars.length; i++) {
            int value = (int) (Math.random() * 26);
            chars[i] = (char) (97 + value);
        }
        return String.valueOf(chars);
    }

}
